!!! written by Cui Qitao, for testing the deep atmosphere baroclinic wave test
!!! refered to Ullrich et.al (2014) A Baroclinic Wave Test Case for Deep and Shallow-atmosphere Dynamical Cores
module deepatm_baro_wave_test_mod

    use const_mod
    use formula_mod
    use latlon_parallel_mod
    use block_mod
    use operators_mod
  
    implicit none
  
    private
  
    public deepatm_baro_wave_test_set_ic
    
    real(r8), parameter :: b_para = 2             ! half-width parameter
    real(r8), parameter :: d0     = radius / 6.0  ! horizontal radius of perturbation domain
    real(r8), parameter :: k_para = 3             ! power used for temperature field
    real(r8), parameter :: p0     = 1.0e5_r8      ! surface pressure
    real(r8), parameter :: t_e0   = 310.0_r8      ! surface equatorial temperature
    real(r8), parameter :: t_p0   = 240.0_r8      ! surface polar temperature
    real(r8), parameter :: u_p    = 1.0_r8        ! perturbed wind amplitude 
    real(r8), parameter :: zt     = 1.5e4_r8      ! top of perturbation domain
    real(r8), parameter :: gamma  = 0.005         ! lapse rate
    real(r8), parameter :: lonc   = pi / 9.0
    real(r8), parameter :: latc   = pi2 / 9.0
    real(r8), parameter :: big_a  = 1.0/gamma
    real(r8), parameter :: t_00   = 0.5 * (t_e0 + t_p0)
    real(r8), parameter :: big_b  = (t_e0 - t_p0)/(t_e0 + t_p0)/t_p0
    real(r8), parameter :: big_c  = 0.5 * (k_para + 2.0) * (t_e0 - t_p0) / (t_e0 * t_p0)
    real(r8), parameter :: big_h  = rd * t_00 / g


    ! real(r8), parameter :: alpha = 0.0_r8
    ! real(r8), parameter :: u0    = 35.0_r8   ! m s-1
    ! real(r8), parameter :: t0    = 288.0_r8  ! K
    ! real(r8), parameter :: gamma = 0.005_r8  ! K m-1
    ! real(r8), parameter :: dt    = 4.8e5_r8  ! K
    ! real(r8), parameter :: eta0  = 0.252
    ! real(r8), parameter :: etat  = 0.2       ! Tropopause level
    ! real(r8), parameter :: lonc  = pi / 9.0
    ! real(r8), parameter :: latc  = pi2 / 9.0
    ! real(r8), parameter :: up    = 1.0       ! m s-1
  
  contains

    pure function zeta(z) result(res)
      real(r8) :: z

      res = 1 - 3*(z/zt)**2 + 2*(z/zt)**3
    end function zeta
    pure function tau_1(r) result(res)
      real(r8), intent(in) :: r

      res = big_a *gamma/t_00 * (exp(gamma*(r-radius)/t_00) - 1) + big_b * (1 - 2*((r-radius)/b_para/big_h)**2) * exp(-((r-radius)/b_para/big_h)**2)
    end function tau_1

    pure function tau_2(r) result(res)
    real(r8), intent(in) :: r

      res = big_c * (1 - 2*((r-radius)/b_para/big_h)**2) * exp(-((r-radius)/b_para/big_h)**2)
    end function tau_2

    pure function tau_1_integral(r) result(res)
      real(r8), intent(in) :: r

      res = big_a * (exp(gamma*(r-radius)/t_00) - 1) + big_b * (r-radius) * exp(-((r-radius)/b_para/big_h)**2)
    end function tau_1_integral

    pure function tau_2_integral(r) result(res)
    real(r8), intent(in) :: r

      res = big_c * (r-radius) * exp(-((r-radius)/b_para/big_h)**2)
    end function tau_2_integral

    ! pure function u_phi_r(phi, r) result(res)
    !   real(r8), intent(in) :: phi
    !   real(r8), intent(in) :: r

    !   res = g/a
    ! end function u_phi_rn

    pure function f_phi_rn(p, phi, r) result(res)
      real(r8), intent(in) :: p
      real(r8), intent(in) :: phi
      real(r8), intent(in) :: r

      res = log(p/p0) + g / rd * tau_1_integral(r) - g / rd * tau_2_integral(r) * ((r*cos(phi)/radius)**k-(k/k+2.0)*(r*cos(phi)/radius)**(k+2))
    end function f_phi_rn

    pure function round_f_round_r(phi, r) result(res)
      real(r8), intent(in) :: phi
      real(r8), intent(in) :: r
        res = big_a * g * gamma / rd / t_00 * exp(gamma/t_00 * (r-radius)) - &
              big_c * g * k_para * cos(phi) *(r-radius) / rd /radius * exp(-((r-radius)/b_para/big_h)**2) * ((r*cos(phi)/radius)**(k-1)-(r*cos(phi)/radius)**(k+1)) + &
              g / rd * (1 - 2*((r-radius)/b_para/big_h)**2) * exp(-((r-radius)/b_para/big_h)**2) * &
              (big_b - big_c((r*cos(phi)/radius)**k-(k/k+2.0)*(r*cos(phi)/radius)**(k+2)))
    end function round_f_round_r

    recursive subroutine solve_r_newton(p, phi, r, iteration_times) !!!tomorow
      real(r8), intent(in) :: p
      real(r8), intent(in) :: phi
      real(r8), intent(inout) :: r              ! first guess value

      integer, intent(inout)  :: iteration_times

      r =  r - f_phi_rn(p, phi, r)/round_f_round_r(phi, r)
      iteration_times = iteration_times - 1
      if (iteration_times == 0) stop
      call solve_r_newton(p, phi, r, iteration_times)

    end subroutine solve_r_newton
  
    subroutine deepatm_baro_wave_test_set_ic(block)
  
      type(block_type), intent(inout), target :: block
  
      real(r8) etav, eta, tbar, ubar, gzbar, sin_lat, cos_lat, lamda, half_lon, d, phi, big_u
      integer i, j, k
      
      real(r8) :: rp !! calc r from p
      real(r8),allocatable :: r_full_3d, r_half_3d
      associate (mesh   => block%mesh            , &
                 u      => block%dstate(1)%u_lon , &
                 v      => block%dstate(1)%v_lat , &
                 mgs    => block%dstate(1)%mgs   , &
                 mg     => block%dstate(1)%mg    , &
                 t      => block%dstate(1)%t     , &
                 pt     => block%dstate(1)%pt    , &
                 gz_lev => block%dstate(1)%gz_lev, &
                 gz     => block%dstate(1)%gz    , &
                 gzs    => block%static%gzs)
      mgs%d = 1.0e5_r8
      v  %d = 0
      allocate(r_full_3d(mesh%full_ids:mesh%full_ide,mesh%full_ids:mesh%full_ids,mesh%full_ids:mesh%full_ids))
      call calc_mg(block, block%dstate(1))
      
      r_full_3d = radius + 10000

      do k = mesh%full_kds, mesh%full_kde
        do j = mesh%full_jds, mesh%full_jde
          phi = mesh%full_lat(j)
          do i = mesh%full_ids, mesh%full_ide
            ! rp = radius + 10000
            call solve_r_newton(mg%d(i,j,k),phi,r_full_3d(i,j,k),5)
          end do
        end do
      end do

      do k = mesh%full_kds, mesh%full_kde
        do j = mesh%full_jds, mesh%full_jde
          phi = mesh%full_lat(j)
          rp = r_full_3d(15,j,k)              ! i like 15
          tbar = (radius/rp)**2 / ( (tau_1(rp) - tau_2(rp) * ((rp*cos(phi)/radius)**k- (k/k+2.0)*(rp*cos(phi)/radius)**(k+2)) ) )
          do i = mesh%full_ids, mesh%full_ide
            !rp = r_full_3d(i,j,k)
            t%d(i,j,k) = tbar 

            pt%d(i,j,k) = modified_potential_temperature(t%d(i,j,k), mg%d(i,j,k), 0.0_r8)
          end do
        end do
      end do
      call fill_halo(pt)

      do k = mesh%full_kds, mesh%full_kde
        do j = mesh%full_jds, mesh%full_jde
          phi   = mesh%full_lat(j)
          rp    = r_full_3d(15,j,k)
          big_u = g/radius * k_para * t%d(15,j,k) *tau_2_integral(rp)*((rp*cos(phi)/radius)**(k-1)-(rp*cos(phi)/radius)**(k+1))
          ubar  = - omega * rp * cos(phi) + sqrt( (omega*rp*cos(phi))**2 + rp* cos(phi) * big_u)
          ! r = radius * acos (sin(latc)*sin(phi)+cos(latc)*cos(phi)*cos)
          do i = mesh%half_ids, mesh%half_ide
            lamda = mesh%half_lon(i)
            d = radius * acos (sin(latc)*sin(phi)+cos(latc)*cos(phi)*cos(lamda - lonc))
            if (d <= d0 ) then
              u%d(i,j,k) = ubar - 16*up/(3*sqrt(3)) * zeta(rp - a) * (cos(pi*d/2/d0))**3 * sin(pi*d/2/d0) * &
                           (-sin(latc)*sin(phi)+cos(latc)*cos(phi)*cos(lamda - lonc)) / sin(d/radius)
              v%d(i,j,k) = 16*up/(3*sqrt(3)) * zeta(rp - a) * (cos(pi*d/2/d0))**3 * sin(pi*d/2/d0) * cos(latc)*sin(lamda-lonc)/sin(d/radius)
            else
            ! rp = r_full_3d(i,j,k)
              u%d(i,j,k) = ubar 
            end if
          end do
        end do
      end do
      call fill_halo(u)
      call fill_halo(v)





      do
        eta = mesh%full_lev(k)
        etav = (eta - eta0) * pi / 2d0
        if (etat <= eta .and. eta <= 1) then
          tbar = t0 * eta**(Rd * gamma / g)
        else
          tbar = t0 * eta**(Rd * gamma / g) + dt * (etat - eta)**5
        end if
        do j = mesh%full_jds, mesh%full_jde
          sin_lat = mesh%full_sin_lat(j)
          cos_lat = mesh%full_cos_lat(j)
          do i = mesh%full_ids, mesh%full_ide
            t%d(i,j,k) = tbar + 3.0d0 / 4.0d0 * eta * pi * u0 / rd * sin(etav) * sqrt(cos(etav)) * (             &
                (-2 * sin_lat**6 * (cos_lat**2 + 1.0d0 / 3.0d0) + 10.0d0 / 63.0d0) * 2 * u0 * cos(etav)**1.5d0 + &
                (8.0d0 / 5.0d0 * cos_lat**3 * (sin_lat**2 + 2.0d0 / 3.0d0) - pi / 4.0d0) * radius * omega        &
              )
            pt%d(i,j,k) = modified_potential_temperature(t%d(i,j,k), mg%d(i,j,k), 0.0_r8)
          end do
        end do
      end do
      call fill_halo(pt)
      









      !!!!!!!!old baro from mr.dong
  
      do k = mesh%full_kds, mesh%full_kde
        eta = mesh%full_lev(k)
        etav = (eta - eta0) * pi05
        do j = mesh%full_jds, mesh%full_jde
          sin_lat = mesh%full_sin_lat(j)
          cos_lat = mesh%full_cos_lat(j)
          do i = mesh%half_ids, mesh%half_ide
            half_lon = mesh%half_lon(i)
            r = 10.0 * acos(sin(latc) * sin_lat + cos(latc) * cos_lat * cos(half_lon - lonc))
            u%d(i,j,k) = u0 * cos(etav)**(1.5d0) * sin(2 * mesh%full_lat(j))**2 + &
                         up * merge(exp(-r**2), 0.0_r8, r**2 <= 50) ! FIXME: Why we set a limit on r?
          end do
        end do
      end do
      call fill_halo(u)
  
      do k = mesh%full_kds, mesh%full_kde
        eta = mesh%full_lev(k)
        etav = (eta - eta0) * pi / 2d0
        if (etat <= eta .and. eta <= 1) then
          tbar = t0 * eta**(Rd * gamma / g)
        else
          tbar = t0 * eta**(Rd * gamma / g) + dt * (etat - eta)**5
        end if
        do j = mesh%full_jds, mesh%full_jde
          sin_lat = mesh%full_sin_lat(j)
          cos_lat = mesh%full_cos_lat(j)
          do i = mesh%full_ids, mesh%full_ide
            t%d(i,j,k) = tbar + 3.0d0 / 4.0d0 * eta * pi * u0 / rd * sin(etav) * sqrt(cos(etav)) * (             &
                (-2 * sin_lat**6 * (cos_lat**2 + 1.0d0 / 3.0d0) + 10.0d0 / 63.0d0) * 2 * u0 * cos(etav)**1.5d0 + &
                (8.0d0 / 5.0d0 * cos_lat**3 * (sin_lat**2 + 2.0d0 / 3.0d0) - pi / 4.0d0) * radius * omega        &
              )
            pt%d(i,j,k) = modified_potential_temperature(t%d(i,j,k), mg%d(i,j,k), 0.0_r8)
          end do
        end do
      end do
      call fill_halo(pt)
  
      do k = mesh%half_kds, mesh%half_kde
        eta = merge(1.0d-12, mesh%half_lev(k), mesh%half_lev(k) == 0)
        etav = (eta - eta0) * pi / 2d0
        if (etat <= eta .and. eta <= 1) then
          gzbar = t0 * g / gamma * (1 - eta**(Rd * gamma / g))
        else
          gzbar = t0 * g / gamma * (1 - eta**(Rd * gamma / g)) - Rd * dt * (       &
              (log(eta / etat) + 137.0d0 / 60.0d0) * etat**5 - 5 * etat**4 * eta + &
              5 * etat**3 * eta**2 - 10.0d0 / 3.0d0 * etat**2 * eta**3 +           &
              5.0d0 / 4.0d0 * etat * eta**4 - 1.0d0 / 5.0d0 * eta**5               &
            )
        end if
        do j = mesh%full_jds, mesh%full_jde
          sin_lat = mesh%full_sin_lat(j)
          cos_lat = mesh%full_cos_lat(j)
          do i = mesh%full_ids, mesh%full_ide
            gz_lev%d(i,j,k) = gzbar + u0 * cos(etav)**1.5d0 * (                                            &
              (-2 * sin_lat**6 * (cos_lat**2 + 1.0d0 / 3.0d0) + 10.0d0 / 63.0d0) * u0 * cos(etav)**1.5d0 + &
              (8.0d0 / 5.0d0 * cos_lat**3 * (sin_lat**2 + 2.0d0 / 3.0d0) - pi / 4.0d0) * radius * omega    &
            )
          end do
        end do
      end do
      call fill_halo(gz_lev)
  
      do k = mesh%full_kds, mesh%full_kde
        eta = mesh%full_lev(k)
        etav = (eta - eta0) * pi / 2d0
        if (etat <= eta .and. eta <= 1) then
          gzbar = t0 * g / gamma * (1 - eta**(Rd * gamma / g))
        else
          gzbar = t0 * g / gamma * (1 - eta**(Rd * gamma / g)) - Rd * dt * (   &
              (log(eta / etat) + 137d0 / 60d0) * etat**5 - 5 * etat**4 * eta + &
              5 * etat**3 * eta**2 - 10d0 / 3d0 * etat**2 * eta**3 +           &
              5d0 / 4d0 * etat * eta**4 - 1d0 / 5d0 * eta**5                   &
            )
        end if
        do j = mesh%full_jds, mesh%full_jde
          sin_lat = mesh%full_sin_lat(j)
          cos_lat = mesh%full_cos_lat(j)
          do i = mesh%full_ids, mesh%full_ide
            gz%d(i,j,k) = gzbar + u0 * cos(etav)**1.5d0 * (                                        &
              (-2 * sin_lat**6 * (cos_lat**2 + 1d0 / 3d0) + 10d0 / 63d0) * u0 * cos(etav)**1.5d0 + &
              (8d0 / 5d0 * cos_lat**3 * (sin_lat**2 + 2d0 / 3d0) - pi / 4d0) * radius * omega      &
            )
          end do
        end do
      end do
      call fill_halo(gz)
  
      etav = (1 - eta0) * pi / 2
      do j = mesh%full_jds, mesh%full_jde
        sin_lat = mesh%full_sin_lat(j)
        cos_lat = mesh%full_cos_lat(j)
        do i = mesh%full_ids, mesh%full_ide
          gzs%d(i,j) = u0 * cos(etav)**1.5d0 * (                                                 &
            (-2 * sin_lat**6 * (cos_lat**2 + 1d0 / 3d0) + 10d0 / 63d0) * u0 * cos(etav)**1.5d0 + &
            (8d0 / 5d0 * cos_lat**3 * (sin_lat**2 + 2d0 / 3d0) - pi / 4d0) * radius * omega      &
          )
        end do
      end do
      call fill_halo(gzs)
      end associate
  
    end subroutine deepatm_baro_wave_test_set_ic
  
  end module deepatm_baro_wave_test_mod
  